#!/usr/bin/env python
# coding=utf-8
from absl import app, flags, logging
import pandas as pd
from gm.api import set_serv_addr, set_token, history, ADJUST_NONE
import pathlib

flags.DEFINE_string('cache_dir', '%s/.cache/quant/stock' % pathlib.Path.home(),
                    'Cache directory')
flags.DEFINE_string('start_date', None, 'Data start date')
flags.DEFINE_string('end_date', None, 'Data end date')
flags.DEFINE_list('order_book_ids', [], 'order_book_ids to read')
flags.DEFINE_enum('freq', '1d', ['1d', '1m'], 'k bar frequency')

FLAGS = flags.FLAGS

frequency = '60s'


def main(argv):
    del argv
    cache_dir = pathlib.Path(FLAGS.cache_dir)
    assert cache_dir.exists()
    set_token('567120cd2f9b2ce8ae3fbd7ee76ed06bf899b512')
    set_serv_addr('192.168.1.122:7001')

    if FLAGS.freq == '1d':
        frequency = '1d'
    else:
        frequency = '60s'
    kdir = cache_dir.joinpath(FLAGS.freq)
    for order_book_id in FLAGS.order_book_ids:
        kpath = kdir.joinpath('%s_%s_%s.csv' % (order_book_id, FLAGS.start_date, FLAGS.end_date))
        if not kpath.exists():
            logging.info('Read %s' % order_book_id)
            df = history(symbol=order_book_id,
                         frequency=frequency,
                         start_time='%s 08:00:00' % FLAGS.start_date,
                         end_time='%s 16:00:00' % FLAGS.end_date,
                         fields='open,high,low,close,volume,amount,eob',
                         adjust=ADJUST_NONE,
                         df=True)
            df.rename(columns={'eob': 'datetime'}, inplace=True)
            df['datetime'] = pd.to_datetime(df.datetime)
            df['order_book_id'] = order_book_id
            df.to_csv(kpath, index=False)


if __name__ == '__main__':
    app.run(main)
